#include "hooks.h"
#include "dfruntime_private.h"
#include "global_queue.h"
#include "page_management.h"
#include "wish_queue.h"
#include "asyncpages.h"
#include <stddef.h>



__thread struct frame_struct **  outptr_tmp;

// Returns bool : True->execute False->abort.
bool hook_before_task() {

    // Get the frame pointer :
    struct frame_struct *fp = dfrt_get_cfp();
    _DEBUG( "Hook before task : %p", fp);
    // Check for the frame pointer acquired :
    bool is_acquired = page_is_acquired( fp );


    if ( !is_acquired ) {
        // Check that page is still wanted :
        if ( !page_is_wanted(fp) ) {
            FATAL( "Frame scheduled in inner queue %p is not wanted nor acquired.",fp);
        }

        // Then postpone by returning false :
        DEBUG("Before task hook says NO.");
        return false;
    }


    // Check for every output frame, if not acquired,
    // allocate the buffers.
    // Remember the ret frames :
    outptr_tmp = calloc( fp->infos->numret, sizeof(struct frame_struct *) );
    memcpy( outptr_tmp, fp->ret, fp->infos->numret * sizeof(struct frame_struct*) );

    for ( int i = 0; i < fp->infos->numret; ++i ) {
        bool is_acquired = page_is_acquired(fp->ret[i]);

        if ( is_acquired ) {
            continue;
        } else {
            DEBUG("Page not arrived, using delayed write.");
            // Allocate a dummy buffer :
            size_t bufsize = 0;
            // Gets the buffer size :
            for ( int j = 0; j < fp->infos->wrdesc[i]->count; ++j) {
                struct wrvec curvec = fp->infos->wrdesc[i]->wrvec[j];
                bufsize = max(bufsize, curvec.offset + curvec.size );
            }

            fp->ret[i] = calloc( bufsize, 1 );
        }

    }

    return true;
}

void hook_after_task() {
    struct frame_struct *fp = dfrt_get_cfp();
    _DEBUG( "Hook after task : %p", fp );
    // Process the writes from the buffers.
    for ( int i = 0; i < fp->infos->numret; ++i ) {
        if ( outptr_tmp[i] == fp->ret[i]) {
            // Then write went straight through
            // We have to free the page though.
            DEBUG( "Will release %p as output frame. ", fp->ret[i] );
            safe_release_page(fp->ret[i]);
            perform_tdecs(fp->ret[i]);

            continue;
        } else {
            // This write will be processed eventually.
            bool do_commits = false;
            perform_tdecs(fp->ret[i]);
	// TODO : check if this critical section deserves to be here...
#pragma omp critical (page_man)
 {
            DEBUG("Element %p added to wishlist.", outptr_tmp[i]);
            add_to_wishlist(outptr_tmp[i], outptr_tmp[i]->infos->framesize,
                            fp->infos->wrdesc[i]->wrvec,
                            fp->infos->wrdesc[i]->count,
                            NULL, 0,
                            -(fp->ret[i]->sc),
                            (char*)fp->ret[i]);
            // Now check if the page arrived in the meantime ...
            if ( page_is_acquired(outptr_tmp[i]) ) {
                // Then commit the changes
                do_commits = true;
            }
}       // End of critical section
         // Commit have acquire/release related guarantee to have
         // enough time to execute, but must be done outside any
         // critical section
         if ( do_commits) {
             wishlist_commit_page(outptr_tmp[i]);
         }

        }
    }

    // We can free the frame :
    safe_release_page(fp);
    // Free the outptr_tmp
    free( outptr_tmp );

    // Output memory info :
    DEBUG("Mem info : a %d w %d", pages_acquired(), pages_wanted());
    // Take care of async page management.
    async_process_all_messages();

}

void hook_would_wait() {
    // Transfer
    _DEBUG( "Hook would wait ");
    int amount = gqueue_reacquire_interbuf();
    if ( amount == 0) {
        amount += gqueue_steal_local(STEAL_AMOUNT);
    }

    DEBUG( "%d has recovered %d tasks locally.", dfrt_get_node_num(), amount );
}

void hook_no_work() {
    DEBUG("Hook no work");
    // Steal local and then global
    int amount = gqueue_reacquire_interbuf();
    if ( amount == 0) {
        amount += gqueue_steal_local(STEAL_AMOUNT);
    }

    DEBUG( "%d has recovered %d tasks locally.", dfrt_get_node_num(), amount );

    // Only thread 0 steals globally.
    if ( amount == 0 ) {
        amount = gqueue_steal_tasks( STEAL_AMOUNT );
        DEBUG( "%d has stolen %d tasks.", dfrt_get_node_num(), amount );
    } else {
        return;
    }

}

void hook_before_local_schedule( struct frame_struct * fp ) {
    // TODO : make the acquire for the main
    // frame non blocking.
    async_request_page( fp, fp->infos->framesize, true);
    for ( int i = 0; i < fp->infos->numret; ++ i ) {
        async_request_page(fp->ret[i], 1, true);
    }


}

void hook_at_page_arrival( struct frame_struct * frame) {
    wishlist_commit_page(frame);
}
